import torch
import torch.nn as nn
from torch.nn import functional as F
from torchvision import datasets, transforms
from datasets.train_dataset import Ms1mDataset
from torch.utils.data import DataLoader
from networks.net import ResNet, InsightFace


# per_batch_size = 10
#
# train_dataset = Ms1mDataset('E:/datasets/faces_emore')
# train_loader = DataLoader(dataset=train_dataset, batch_size=per_batch_size, shuffle=True, num_workers=10)

# veri_dataset = VerificationDataset('E:/datasets/faces_emore')
# veri_loader = DataLoader(dataset=veri_dataset, batch_size=per_batch_size, shuffle=False, num_workers=10)

class InsightFace1(nn.Module):
    def __init__(self, num_class, scale=64, act_type='relu', m=(1.0, 0.0, 0.0)):
        super(InsightFace1, self).__init__()
        self.m = m
        self.scale = scale
        self.net = ResNet(act_type=act_type)
        self.cos = nn.Linear(512, num_class, bias=False)
        self.num_class = num_class
        for _layer in self.modules():
            if isinstance(_layer, nn.Conv2d) or isinstance(_layer, nn.Linear):
                nn.init.xavier_normal_(_layer.weight, 2 ** 0.5)
    def forward(self, x, label):
        print(x.shape)
        x = self.net(x)
        print(x.shape)
        x_net = F.normalize(x, dim=1)
        weight_normalize = F.normalize(self.cos.weight, dim=1)
        x_cos = F.linear(x_net, weight_normalize, self.cos.bias)
        print(x_cos.shape)
        onehot_label = torch.zeros_like(x_cos).scatter_(1, label, 1)
        x_cor = torch.cos(torch.acos(x_cos)*self.m[0]+self.m[1])-self.m[2]
        print(x_cor.shape)
        x = onehot_label*x_cor + (1-onehot_label)*x_cos
        x = self.scale*x
        x = F.log_softmax(x, dim=1)
        return x, x_cos, x_net

net = InsightFace(100)
print(net)

a = torch.Tensor(10,3,224,224)
print(a.shape)
label = torch.randint(0,100,(10,1))
b = net(a, label)
print(b[0].shape, b[1].shape, b[2].shape)

